#include "sepsampler.h"

// main sampler structure
void sep_add_sampler(sepsampler *sptr, const char *sampler, 
		     sepsys sys, int lvec, ...){
  va_list args;

  va_start(args, lvec); 
  
  if ( strcmp(sampler, "sacf")==0 ){
    if ( sptr->flag_sacf != 1 ){
      double tsample  = va_arg(args, double);
      sptr->sacf = sep_sacf_init(lvec, tsample, sys.dt);
      sptr->flag_sacf = 1;
    }
  }
  else if ( strcmp(sampler, "vacf")==0 ){
    if ( sptr->flag_vacf != 1 ){
      double tsample  = va_arg(args, double);
      sptr->vacf = sep_vacf_init(lvec, tsample, sys);
      sptr->flag_vacf = 1;
    }
  }
  else if ( strcmp(sampler, "msacf")==0 ){
    if ( sptr->flag_msacf != 1 ){
      
      if ( sptr->molptr == NULL ) 
	sep_error("%s at line %d: molpointer not initialized", 
		  __func__, __LINE__);

      double tsample  = va_arg(args, double);
      sptr->msacf = sep_msacf_init(lvec, tsample, sys.dt);
      sptr->flag_msacf = 1;
    }
  }
  else if ( strcmp(sampler, "dfluxacf")==0 ){

    sep_warning("dfluxacf sampler is in a poor state - check code please!");

    if ( sptr->flag_dfluxacf != 1 ){
      
      if ( sptr->molptr == NULL ) 
	sep_error("%s at line %d: molpointer not initialized", 
		  __func__, __LINE__);

      double tsample  = va_arg(args, double);
      sptr->dfluxacf = sep_dflux_init(lvec, tsample, sys.dt);
      sptr->flag_dfluxacf = 1;
    }
  }
  else if ( strcmp(sampler, "gh")==0 ){
    if ( sptr->flag_gh != 1 ){
      double tsample  = va_arg(args, double);
      int nwave = va_arg(args, int);
      sptr->gh = sep_gh_init(lvec, tsample, sys.dt, nwave, sys.length[1]);
      sptr->flag_gh = 1;
    }
  }
  else if ( strcmp(sampler, "mgh")==0 ){
    if ( sptr->flag_mgh != 1 ){
      if ( sptr->molptr == NULL ) 
	sep_error("%s at line %d: molpointer not initialized", 
		  __func__, __LINE__);

      double tsample  = va_arg(args, double);
      int nwave = va_arg(args, int);
      sptr->mgh = sep_mgh_init(lvec, tsample, sys.dt, nwave, sys.length[1]);
      sptr->flag_mgh = 1;
    }
  }
  else if ( strcmp(sampler, "profs")==0 ){
    if ( sptr->flag_profs != 1 ){
      char type = va_arg(args, int);
      int isample = va_arg(args, int);
      sptr->profs = sep_profs_init(type, lvec, isample);
      sptr->flag_profs = 1;
    }
  }
  else if ( strcmp(sampler, "mprofs")==0 ){
    if ( sptr->flag_mprofs != 1 ){

      char type = va_arg(args, int);
      int isample = va_arg(args, int);
      int dir = va_arg(args, int);
      int dirvel = va_arg(args, int);
      int diramom = va_arg(args, int);

      sptr->mprofs = sep_mprofs_init(type, lvec, isample, dir, dirvel, diramom);
      sptr->flag_mprofs = 1;
    }
  }
  else if ( strcmp(sampler, "mcacf")==0 ){
    if ( sptr->flag_mcacf != 1 ){
      
      if ( sptr->molptr == NULL ) 
	sep_error("%s at line %d: molpointer not initialized", 
		  __func__, __LINE__);

      double tsample  = va_arg(args, double);
      sptr->mcacf = sep_mcacf_init(lvec, tsample, sys.dt);
      sptr->flag_mcacf = 1;
    }
  }
  else if ( strcmp(sampler, "mvacf")==0 ){
    if ( sptr->flag_mvacf != 1 ){
      double tsample  = va_arg(args, double);
      sptr->mvacf = sep_mvacf_init(lvec, tsample, sys);
      sptr->flag_mvacf = 1;
    }
  }
   else if ( strcmp(sampler, "mavacf")==0 ){
    if ( sptr->flag_mavacf != 1 ){
      double tsample  = va_arg(args, double);
      sptr->mavacf = sep_mavacf_init(lvec, tsample, sys);
      sptr->flag_mavacf = 1;
    }
  }
  else {
    sep_error("%s and at line %d: Sampler %s is not recognized", 
	      __func__, __LINE__, sampler);
  }
  
  va_end(args);
}

void sep_add_mol_sampler(sepsampler *sptr, sepmol *mols){
  
  sptr->molptr = mols;

}

void sep_close_sampler(sepsampler *ptr){
  
  if ( ptr->flag_sacf == 1 )    sep_sacf_close(ptr->sacf);
  if ( ptr->flag_vacf == 1 )    sep_vacf_close(ptr->vacf);
  if ( ptr->flag_msacf == 1 )   sep_msacf_close(ptr->msacf);
  if ( ptr->flag_gh == 1 )      sep_gh_close(ptr->gh);
  if ( ptr->flag_mgh == 1 )     sep_mgh_close(ptr->mgh);
  if ( ptr->flag_profs == 1 )   sep_profs_close(ptr->profs);
  if ( ptr->flag_mprofs == 1 )  sep_mprofs_close(ptr->mprofs);
  if ( ptr->flag_mcacf == 1 )   sep_mcacf_close(ptr->mcacf);
  if ( ptr->flag_mvacf == 1 )   sep_mvacf_close(ptr->mvacf);
  if ( ptr->flag_mavacf == 1 )  sep_mavacf_close(ptr->mavacf);
  if ( ptr->flag_dfluxacf == 1 )  sep_dflux_close(ptr->dfluxacf);

}

void sep_sample(seppart *pptr, sepsampler *sptr, sepret *ret, sepsys sys, 
		unsigned n){
  
  if ( sptr->flag_sacf==1 && n%sptr->sacf->isample == 0 )
    sep_sacf_sample(sptr->sacf, ret, sys);
  
  if ( sptr->flag_vacf==1 && n%sptr->vacf->isample == 0 )
    sep_vacf_sample(pptr, sptr->vacf, sys);

  if ( sptr->flag_msacf==1 && n%sptr->msacf->isample == 0 )
    sep_msacf_sample(sptr->msacf, pptr, sptr->molptr, ret, sys);

  if ( sptr->flag_gh==1 && n%sptr->gh->isample == 0 )
    sep_gh_sampler(pptr, sptr->gh, sys);

  if ( sptr->flag_mgh==1 && n%sptr->mgh->isample == 0 )
    sep_mgh_sampler(pptr, sptr->molptr, sptr->mgh, sys);

  if ( sptr->flag_profs==1 && n%sptr->profs->isample == 0 )
    sep_profs_sampler(pptr, sptr->profs, sys);

  if ( sptr->flag_mprofs==1 && n%sptr->mprofs->isample == 0 )
    sep_mprofs_sampler(pptr, sptr->molptr, sptr->mprofs, sys);

  if ( sptr->flag_mcacf==1 && n%sptr->mcacf->isample == 0 )
    sep_mcacf_sample(sptr->mcacf, pptr, sptr->molptr, ret, sys);

  if ( sptr->flag_mvacf==1 && n%sptr->mvacf->isample == 0 )
    sep_mvacf_sample(pptr, sptr->mvacf, sptr->molptr, sys);  

  if ( sptr->flag_mavacf==1 && n%sptr->mavacf->isample == 0 )
    sep_mavacf_sample(pptr, sptr->mavacf, sptr->molptr, sys);  

  if ( sptr->flag_dfluxacf==1 && n%sptr->dfluxacf->isample == 0 )
    sep_dflux_sample(pptr, sptr->dfluxacf, sptr->molptr, sys);  
    
}

sepsampler sep_init_sampler(void){
  
  sepsampler a;

  a.molptr = NULL;

  a.flag_sacf = 0;  a.sacf = NULL;
  a.flag_msacf = 0;  a.msacf = NULL;

  a.flag_vacf = 0;  a.vacf = NULL;
  a.flag_mvacf = 0;  a.mvacf = NULL;
  a.flag_mavacf = 0;  a.mavacf = NULL;

  a.flag_gh = 0; a.gh = NULL;
  a.flag_mgh = 0; a.mgh = NULL;

  a.flag_profs = 0; a.profs = NULL;
  a.flag_mprofs = 0; a.mprofs = NULL;

  a.flag_mcacf = 0; a.mcacf = NULL;

  a.flag_dfluxacf = 0; a.dfluxacf = NULL;

  return a;
}

/************************************************
 *
 * Function definitions for the individual samplers
 *
 *************************************************/

// Stress ACF
sepsacf *sep_sacf_init(int lvec, double tsample, double dt){
  
  sepsacf *ptr = malloc(sizeof(sepsacf));
  if ( ptr==NULL )
    sep_error("%s at line %d: Couldn't allocate memory", __func__, __LINE__);
  
  ptr->dt   = dt;
  ptr->lvec = lvec;
  ptr->i    = 0; 
  
  ptr->nsample = 0;

  ptr->dtsample = tsample/lvec;
  ptr->isample = (int)(ptr->dtsample/dt);
  if ( ptr->isample < 1 ) 
    sep_error("%s at line %d: isample is too small - CHECK lvec argument",
	      __func__, __LINE__);
  
  ptr->sacf = sep_vector(lvec);
  ptr->stress = sep_matrix(lvec, 3);

  return ptr;
}


void sep_sacf_sample(sepsacf *sptr, sepret *ret, sepsys sys){

  sep_pressure_tensor(ret, sys);

  int index = sptr->i;

  sptr->stress[index][0] = -ret->P[0][1];
  sptr->stress[index][1] = -ret->P[0][2];
  sptr->stress[index][2] = -ret->P[1][2];

  (sptr->i)++;

  if ( sptr->i == sptr->lvec ){
     
    for ( int k=0; k<3; k++ ){
      for ( unsigned n=0; n<sptr->lvec; n++ ){
        for ( unsigned nn=0; nn<sptr->lvec-n; nn++ ){
          sptr->sacf[n] += sptr->stress[nn][k]*sptr->stress[n+nn][k];
        }
      }
    }

    (sptr->nsample)++;

    FILE *fout1 = fopen("sacf.dat", "w");
   
    for ( unsigned n=0; n<sptr->lvec; n++ ){
      double t   = (n+0.5)*sptr->dtsample;
      double fac = sys.volume/(3*(sptr->lvec-n)*sptr->nsample);
      fprintf(fout1, "%f %f\n", t, sptr->sacf[n]*fac);
    }

    fclose(fout1); 
    
    sptr->i = 0;
    
  }

}

void sep_sacf_close(sepsacf *ptr){
  
  free(ptr->sacf);
  sep_free_matrix(ptr->stress, ptr->lvec);
  
  free(ptr);

}


// Velocity ACF
sepvacf *sep_vacf_init(int lvec, double tsample, sepsys sys){
  
  sepvacf *ptr = malloc(sizeof(sepvacf));
  if ( ptr==NULL )
    sep_error("%s at line %d: Couldn't allocate memory", __func__, __LINE__);
  
  ptr->dt   = sys.dt;
  ptr->lvec = lvec;
  ptr->i    = 0; 
  
  ptr->nsample = 0;
  
  ptr->dtsample = tsample/lvec;
  ptr->isample = (int)(ptr->dtsample/sys.dt);
  if ( ptr->isample < 1 ) 
    sep_error("%s at line %d: isample is too small - CHECK lvec argument",
	      __func__, __LINE__);

  ptr->vacf = sep_vector(lvec);
  ptr->vels = sep_matrix(lvec, sys.npart);
  
  return ptr;
}

 
void sep_vacf_sample(seppart *ptr, sepvacf *vptr, sepsys sys){
  
  int index = vptr->i;
  for ( int i=0; i<sys.npart; i++ )
    vptr->vels[index][i] = ptr[i].v[0];
  
  (vptr->i)++;
  
  if ( vptr->i == vptr->lvec ){
    
    for ( int i=0; i<sys.npart; i++ )
      for ( unsigned k=0; k<vptr->lvec; k++ )
        for ( unsigned kk=0; kk<vptr->lvec-k; kk++ )
          vptr->vacf[k] += vptr->vels[kk][i]*vptr->vels[k+kk][i];   
  
    (vptr->nsample)++;

    FILE *fout = fopen("vacf.dat", "w");
    for ( unsigned k=0; k<vptr->lvec; k++ ){
      double t   = (k+0.5)*vptr->dtsample;
      double fac = 1.0/((vptr->lvec-k)*sys.npart*vptr->nsample);
      fprintf(fout, "%f %f\n", t, vptr->vacf[k]*fac);
    }
    fclose(fout);

    vptr->i=0;
    
  }

}

void sep_vacf_close(sepvacf *ptr){

  free(ptr->vacf);
  sep_free_matrix(ptr->vels, ptr->lvec);
  
  free(ptr);

}


// Molecular stress ACF
sepmsacf *sep_msacf_init(int lvec, double tsample, double dt){
  
  sepmsacf *ptr = malloc(sizeof(sepmsacf));
  if ( ptr==NULL )
    sep_error("%s at line %d: Couldn't allocate memory", __func__, __LINE__);
  
  ptr->dt   = dt;
  ptr->lvec = lvec;
  ptr->i    = 0; 
  
  ptr->nsample = 0;

  ptr->dtsample = tsample/lvec;
  ptr->isample = (int)(ptr->dtsample/dt);
  if ( ptr->isample < 1 ) 
    sep_error("%s at line %d: isample is too small - CHECK lvec argument",
	      __func__, __LINE__);


  ptr->sacf = sep_matrix(lvec, 2);
  ptr->sstress = sep_matrix(lvec, 3);
  ptr->astress = sep_matrix(lvec, 3);

  return ptr;
}


void sep_msacf_sample(sepmsacf *sptr, sepatom *atoms, sepmol *mols, 
		      sepret *ret, sepsys sys){

  sep_mol_pressure_tensor(atoms, mols, ret, sys);

  int index = sptr->i;

  sptr->sstress[index][0] = 0.5*(ret->P_mol[0][1]+ret->P_mol[1][0]);
  sptr->sstress[index][1] = 0.5*(ret->P_mol[0][2]+ret->P_mol[2][0]);
  sptr->sstress[index][2] = 0.5*(ret->P_mol[1][2]+ret->P_mol[2][1]);

  sptr->astress[index][0] = 0.5*(ret->P_mol[0][1]-ret->P_mol[1][0]);
  sptr->astress[index][1] = 0.5*(ret->P_mol[0][2]-ret->P_mol[2][0]);
  sptr->astress[index][2] = 0.5*(ret->P_mol[1][2]-ret->P_mol[2][1]);

  (sptr->i)++;

  if ( sptr->i == sptr->lvec ){
     
    for ( int k=0; k<3; k++ ){
      for ( unsigned n=0; n<sptr->lvec; n++ ){
        for ( unsigned nn=0; nn<sptr->lvec-n; nn++ ){
          sptr->sacf[n][0] += sptr->sstress[nn][k]*sptr->sstress[n+nn][k];
	  sptr->sacf[n][1] += sptr->astress[nn][k]*sptr->astress[n+nn][k];
        }
      }
    }

    (sptr->nsample)++;

    FILE *fout1 = fopen("msacf.dat", "w");
   
    for ( unsigned n=0; n<sptr->lvec; n++ ){
      double t   = (n+0.5)*sptr->dtsample;
      double fac = sys.volume/(3*(sptr->lvec-n)*sptr->nsample);
      fprintf(fout1, "%f %f %f\n", t, sptr->sacf[n][0]*fac, 
	      sptr->sacf[n][1]*fac);
    }

    fclose(fout1); 
    
    sptr->i = 0;
    
  }

}

void sep_msacf_close(sepmsacf *ptr){
  
 
  sep_free_matrix(ptr->sacf, ptr->lvec);
  sep_free_matrix(ptr->sstress, ptr->lvec);
  sep_free_matrix(ptr->astress, ptr->lvec);
  
  free(ptr);

}


// Dipole flux tensor
sepdfluxacf *sep_dflux_init(int lvec, double tsample, double dt){
  
  sepdfluxacf *ptr = malloc(sizeof(sepmsacf));
  if ( ptr==NULL )
    sep_error("%s at line %d: Couldn't allocate memory", __func__, __LINE__);
  
  ptr->dt   = dt;
  ptr->lvec = lvec;
  ptr->i    = 0; 
  
  ptr->nsample = 0;

  ptr->dtsample = tsample/lvec;
  ptr->isample = (int)(ptr->dtsample/dt);
  if ( ptr->isample < 1 ) 
    sep_error("%s at line %d: isample is too small - CHECK lvec argument",
	      __func__, __LINE__);


  ptr->acf = sep_matrix(lvec, 3);
  ptr->dflux = sep_matrix(lvec, 7);

  return ptr;
}


void sep_dflux_sample(sepatom *atoms, sepdfluxacf *sptr, sepmol *mols, sepsys sys){
  

  double **R = sep_matrix(3,3);
  double **Ros = sep_matrix(3,3);
  double **Ra = sep_matrix(3,3);
  double r;  

  sep_eval_dipole_flux_tensor(R, atoms, mols, sys);
  sep_decompose_tensor(&r, Ros, Ra, R, 3);
  
  int index = sptr->i;

  sptr->dflux[index][0] = Ros[0][1]; 
  sptr->dflux[index][1] = Ros[0][2];
  sptr->dflux[index][2] = Ros[1][2];

  sptr->dflux[index][3] = Ra[1][2];
  sptr->dflux[index][4] = Ra[2][0]; 
  sptr->dflux[index][5] = Ra[0][1];

  sptr->dflux[index][6] = r;

  (sptr->i)++;

  if ( sptr->i == sptr->lvec ){
     
    for ( unsigned n=0; n<sptr->lvec; n++ ){
      for ( unsigned nn=0; nn<sptr->lvec-n; nn++ ){
	sptr->acf[n][0] += sptr->dflux[nn][0]*sptr->dflux[n+nn][0] + 
	  sptr->dflux[nn][1]*sptr->dflux[n+nn][1] + 
	  sptr->dflux[nn][2]*sptr->dflux[n+nn][2] ; 

 	sptr->acf[n][1] += sptr->dflux[nn][3]*sptr->dflux[n+nn][3] + 
	  sptr->dflux[nn][4]*sptr->dflux[n+nn][4] + 
	  sptr->dflux[nn][5]*sptr->dflux[n+nn][5] ; 

	sptr->acf[n][2] += sptr->dflux[nn][6]*sptr->dflux[n+nn][6];
      }
    }
    (sptr->nsample)++;

    FILE *fout1 = fopen("dflux.dat", "w");
   
    for ( unsigned n=0; n<sptr->lvec; n++ ){
      double t   = (n+0.5)*sptr->dtsample;
      double fac = sys.volume/(3*(sptr->lvec-n)*sptr->nsample);
      fprintf(fout1, "%f %f %f %f\n", t, sptr->acf[n][0]*fac, 
	      sptr->acf[n][1]*fac, sptr->acf[n][2]*fac*3);   
    }

    fclose(fout1); 
    
    sptr->i = 0;
    
  }

  sep_free_matrix(R, 3);
  sep_free_matrix(Ros, 3);
  sep_free_matrix(Ra, 3);
}

void sep_dflux_close(sepdfluxacf *ptr){
  
 
  sep_free_matrix(ptr->acf, ptr->lvec);
  sep_free_matrix(ptr->dflux, ptr->lvec);

  free(ptr);

}




// Gen. hydro sampler (atomic)
sepgh *sep_gh_init(int lvec, double tsample, double dt, int nwave, 
		   double Ly){
  
  sepgh *ptr = malloc(sizeof(sepgh));
  if ( ptr==NULL )
    sep_error("%s at line %d: Couldn't allocate memory", __func__, __LINE__);
  
  ptr->dt   = dt;
  ptr->lvec = lvec;
  ptr->i    = 0; 
  
  ptr->nsample = 0;

  ptr->dtsample = tsample/lvec;
  ptr->isample = (int)(ptr->dtsample/dt);
  if ( ptr->isample < 1 ) 
    sep_error("%s at line %d: isample is too small - CHECK lvec argument",
	      __func__, __LINE__);

  ptr->nwave = nwave;
  ptr->k = sep_vector(nwave);
  FILE *fout = fopen("k.dat", "w");
  if ( fout == NULL )
    sep_error("%s at %d: Couldn't open file k.dat");

  for ( int n=1; n<=nwave; n++ ){
    ptr->k[n-1] = 2*SEP_PI*n/Ly;
    fprintf(fout, "%f\n", ptr->k[n-1]);
  }

  fclose(fout);

  ptr->fk_tv  = sep_complex_matrix(lvec, nwave);
  ptr->fkm_tv = sep_complex_matrix(lvec, nwave);
  ptr->fk_lv  = sep_complex_matrix(lvec, nwave);
  ptr->fkm_lv = sep_complex_matrix(lvec, nwave);
  ptr->fk_scatt  = sep_complex_matrix(lvec, nwave);
  ptr->fkm_scatt = sep_complex_matrix(lvec, nwave);

  ptr->fk_cdens  = sep_complex_matrix(lvec, nwave);
  ptr->fkm_cdens = sep_complex_matrix(lvec, nwave);

  ptr->c_tv     = sep_complex_matrix(lvec, nwave);
  ptr->c_lv     = sep_complex_matrix(lvec, nwave);
  ptr->c_scatt  = sep_complex_matrix(lvec, nwave);
  ptr->c_cdens  = sep_complex_matrix(lvec, nwave);

  
  return ptr;
}

void sep_gh_close(sepgh *ptr){
  int lvec = ptr->lvec;

  sep_free_complex_matrix(ptr->fk_tv, lvec);
  sep_free_complex_matrix(ptr->fkm_tv, lvec);
  sep_free_complex_matrix(ptr->fk_lv, lvec);
  sep_free_complex_matrix(ptr->fkm_lv, lvec);
  sep_free_complex_matrix(ptr->fk_scatt, lvec);
  sep_free_complex_matrix(ptr->fkm_scatt, lvec);
  sep_free_complex_matrix(ptr->fk_cdens, lvec);
  sep_free_complex_matrix(ptr->fkm_cdens, lvec);

  sep_free_complex_matrix(ptr->c_lv, lvec);
  sep_free_complex_matrix(ptr->c_scatt, lvec);
  sep_free_complex_matrix(ptr->c_tv, lvec);
  sep_free_complex_matrix(ptr->c_cdens, lvec);

}


void sep_gh_sampler(sepatom *ptr, sepgh *gh, sepsys sys){

  int index = gh->i;

  sep_eval_xtrue(ptr, sys);

  for ( unsigned n=0; n<gh->nwave; n++ ){
    
    gh->fk_tv[index][n]    = gh->fkm_tv[index][n]    = 0.0 + 0.0*I;
    gh->fk_lv[index][n]    = gh->fkm_lv[index][n]    = 0.0 + 0.0*I;
    gh->fk_scatt[index][n] = gh->fkm_scatt[index][n] = 0.0 + 0.0*I;
    gh->fk_cdens[index][n] = gh->fkm_cdens[index][n] = 0.0 + 0.0*I;

    for ( int m=0; m<sys.npart; m++ ){
      if ( ptr[n].type== 'A'){
	complex double kfac   = cexp(I*gh->k[n]*ptr[m].xtrue[1]);
	complex double kfac_m = cexp(-I*gh->k[n]*ptr[m].xtrue[1]);

	//complex double kfac   = cexp(I*gh->k[n]*ptr[m].x[1]);
	//complex double kfac_m = cexp(-I*gh->k[n]*ptr[m].x[1]);
	
	gh->fk_tv[index][n]   +=  ptr[m].v[2]*kfac;
	gh->fkm_tv[index][n]  +=  ptr[m].v[2]*kfac_m;
	
	gh->fk_lv[index][n]   +=  ptr[m].v[1]*kfac;
	gh->fkm_lv[index][n]  +=  ptr[m].v[1]*kfac_m;

	gh->fk_scatt[index][n]   +=  kfac;
	gh->fkm_scatt[index][n]  +=  kfac_m;
	
	gh->fk_cdens[index][n]   +=  ptr[m].z*kfac;
	gh->fkm_cdens[index][n]  +=  ptr[m].z*kfac_m;
	
      }
    }
  }

  (gh->i)++;

  if ( gh->i == gh->lvec ){
    for ( unsigned k=0; k<gh->nwave; k++ ){
      for ( unsigned n=0; n<gh->lvec; n++ ){
	for ( unsigned nn=0; nn<gh->lvec-n; nn++ ){
	  gh->c_tv[n][k]    += gh->fkm_tv[nn][k]*gh->fk_tv[n+nn][k];
	  gh->c_lv[n][k]    += gh->fkm_lv[nn][k]*gh->fk_lv[n+nn][k];
	  gh->c_scatt[n][k] += gh->fk_scatt[nn][k]*gh->fkm_scatt[n+nn][k];
	  gh->c_cdens[n][k] += gh->fk_cdens[nn][k]*gh->fkm_cdens[n+nn][k];
	}
      }
    }

    (gh->nsample)++;
    
    FILE *fout_tv = fopen("tvacf.dat", "w");
    FILE *fout_lv = fopen("lvacf.dat", "w");
    FILE *fout_scatt = fopen("scatt.dat", "w");
    FILE *fout_cdens = fopen("rhozacf.dat", "w");

    for ( unsigned n=0; n<gh->lvec; n++ ){
      double t   = (n+0.5)*gh->dtsample;
      
      fprintf(fout_tv, "%f ", t);
      fprintf(fout_lv, "%f ", t);
      fprintf(fout_scatt, "%f ", t);
      fprintf(fout_cdens, "%f ", t);
      
      double fac = 1.0/(gh->nsample*sys.volume*(gh->lvec-n));
      
      for ( unsigned k=0; k<gh->nwave; k++ ){
	fprintf(fout_tv, "%f ", creal(gh->c_tv[n][k]*fac));
	fprintf(fout_lv, "%f ", creal(gh->c_lv[n][k]*fac));
	fprintf(fout_scatt, "%f ", creal(gh->c_scatt[n][k]*fac));
	fprintf(fout_cdens, "%f %f ", creal(gh->c_cdens[n][k]*fac), 
		cimag(gh->c_cdens[n][k]*fac));

      }
      fprintf(fout_tv, "\n"); 
      fprintf(fout_lv, "\n"); 
      fprintf(fout_scatt, "\n");
      fprintf(fout_cdens, "\n");

    }
	
    fclose(fout_tv);
    fclose(fout_lv);
    fclose(fout_scatt);
    fclose(fout_cdens);
    
    gh->i = 0;
  }

  
}


// Molecular gen. hydrodynamics sampler
sepmgh *sep_mgh_init(int lvec, double tsample, double dt, int nwave, 
		    double Ly){
  
  sepmgh *ptr = malloc(sizeof(sepmgh));
  if ( ptr==NULL )
    sep_error("%s at line %d: Couldn't allocate memory", __func__, __LINE__);
  
  ptr->dt   = dt;
  ptr->lvec = lvec;
  ptr->i    = 0; 
  
  ptr->nsample = 0;

  ptr->dtsample = tsample/lvec;
  ptr->isample = (int)(ptr->dtsample/dt);
  if ( ptr->isample < 1 ) 
    sep_error("%s at line %d: isample is too small - CHECK lvec argument",
	      __func__, __LINE__);

  ptr->nwave = nwave;
  ptr->k = sep_vector(nwave);
  FILE *fout = fopen("mk.dat", "w");
  if ( fout == NULL )
    sep_error("%s at %d: Couldn't open file k.dat");

  for ( int n=1; n<=nwave; n++ ){
    ptr->k[n-1] = 2*SEP_PI*n/Ly;
    fprintf(fout, "%f\n", ptr->k[n-1]);
  }

  fclose(fout);

  ptr->fk_tv  = sep_complex_matrix(lvec, nwave);
  ptr->fkm_tv = sep_complex_matrix(lvec, nwave);

  ptr->fk_lv  = sep_complex_matrix(lvec, nwave);
  ptr->fkm_lv = sep_complex_matrix(lvec, nwave);

  ptr->fk_scatt  = sep_complex_matrix(lvec, nwave);
  ptr->fkm_scatt = sep_complex_matrix(lvec, nwave);

  ptr->fk_tav  = sep_complex_matrix(lvec, nwave);
  ptr->fkm_tav = sep_complex_matrix(lvec, nwave);

  ptr->fk_lav  = sep_complex_matrix(lvec, nwave);
  ptr->fkm_lav = sep_complex_matrix(lvec, nwave);

  ptr->fk_vav  = sep_complex_matrix(lvec, nwave);
  ptr->fkm_vav = sep_complex_matrix(lvec, nwave);
 
  ptr->fk_tec  = sep_complex_matrix(lvec, nwave);
  ptr->fkm_tec = sep_complex_matrix(lvec, nwave);

  ptr->fk_lec  = sep_complex_matrix(lvec, nwave);
  ptr->fkm_lec = sep_complex_matrix(lvec, nwave);
 
  ptr->fk_cdens  = sep_complex_matrix(lvec, nwave);
  ptr->fkm_cdens = sep_complex_matrix(lvec, nwave);
  
  ptr->c_tv     = sep_complex_matrix(lvec, nwave);
  ptr->c_lv     = sep_complex_matrix(lvec, nwave);
  ptr->c_scatt  = sep_complex_matrix(lvec, nwave);
  ptr->c_tav    = sep_complex_matrix(lvec, nwave);
  ptr->c_lav    = sep_complex_matrix(lvec, nwave);
  ptr->c_vav    = sep_complex_matrix(lvec, nwave);

  ptr->c_tec    = sep_complex_matrix(lvec, nwave);
  ptr->c_lec    = sep_complex_matrix(lvec, nwave);
  ptr->c_cdens  = sep_complex_matrix(lvec, nwave);
  ptr->c_tcav   = sep_complex_matrix(lvec, nwave);
  ptr->c_lcav   = sep_complex_matrix(lvec, nwave);

  ptr->c_cdensav   = sep_complex_matrix(lvec, nwave);

  return ptr;
}

void sep_mgh_close(sepmgh *ptr){
  int lvec = ptr->lvec;

  sep_free_complex_matrix(ptr->fk_tv, lvec);
  sep_free_complex_matrix(ptr->fkm_tv, lvec);
  sep_free_complex_matrix(ptr->fk_lv, lvec);
  sep_free_complex_matrix(ptr->fkm_lv, lvec);
  sep_free_complex_matrix(ptr->fk_scatt, lvec);
  sep_free_complex_matrix(ptr->fkm_scatt, lvec);
  sep_free_complex_matrix(ptr->fk_tav, lvec);
  sep_free_complex_matrix(ptr->fkm_tav, lvec);
  sep_free_complex_matrix(ptr->fk_lav, lvec);
  sep_free_complex_matrix(ptr->fkm_lav, lvec);
  sep_free_complex_matrix(ptr->fk_vav, lvec);
  sep_free_complex_matrix(ptr->fkm_vav, lvec);
  sep_free_complex_matrix(ptr->fk_tec, lvec);
  sep_free_complex_matrix(ptr->fkm_tec, lvec);
  sep_free_complex_matrix(ptr->fk_lec, lvec);
  sep_free_complex_matrix(ptr->fkm_lec, lvec);
  sep_free_complex_matrix(ptr->fk_cdens, lvec);
  sep_free_complex_matrix(ptr->fkm_cdens, lvec);
 
  sep_free_complex_matrix(ptr->c_lv, lvec);
  sep_free_complex_matrix(ptr->c_scatt, lvec);
  sep_free_complex_matrix(ptr->c_tv, lvec);
  sep_free_complex_matrix(ptr->c_tav, lvec);
  sep_free_complex_matrix(ptr->c_lav, lvec);
  sep_free_complex_matrix(ptr->c_vav, lvec);

  sep_free_complex_matrix(ptr->c_tec, lvec);
  sep_free_complex_matrix(ptr->c_lec, lvec);
  sep_free_complex_matrix(ptr->c_cdens, lvec);
  sep_free_complex_matrix(ptr->c_tcav, lvec);
  sep_free_complex_matrix(ptr->c_lcav, lvec);

  sep_free_complex_matrix(ptr->c_cdensav, lvec);
}

// Only works for single component uniaxial molecules
// See Hansen et al. Langmuir, .... 
void sep_mgh_sampler(sepatom *ptr, sepmol *mols, sepmgh *mgh, sepsys sys){


  sep_mol_spin(ptr, mols, sys, 1); // Here cm is also evaluated
  sep_mol_velcm(ptr, mols, sys);
  sep_mol_eval_xtrue(ptr, mols, sys);  

  complex double kfac, kfac_m;
  unsigned n, m;
  
  int index = mgh->i;
 
  const double mass = mols[0].m;
  double *molcharge = sep_vector(sys.molptr->num_mols);
  double tot_mass = 0.0;
  for ( m=0; m<sys.molptr->num_mols; m++ ){
   
    tot_mass += mols[m].m;
    
    for ( unsigned n=0; n<mols[n].nuau; n++ )
      molcharge[m] += ptr[mols[m].index[n]].z;
  }
  
  const double mdens = tot_mass/sys.volume;

  for ( n=0; n<mgh->nwave; n++ ){
    
    mgh->fk_tv[index][n]    = mgh->fkm_tv[index][n]    = 0.0 + 0.0*I;
    mgh->fk_lv[index][n]    = mgh->fkm_lv[index][n]    = 0.0 + 0.0*I;
    mgh->fk_scatt[index][n] = mgh->fkm_scatt[index][n] = 0.0 + 0.0*I;
    mgh->fk_tav[index][n]    = mgh->fkm_tav[index][n]  = 0.0 + 0.0*I;
    mgh->fk_lav[index][n]    = mgh->fkm_lav[index][n]  = 0.0 + 0.0*I;
    mgh->fk_vav[index][n]    = mgh->fkm_vav[index][n]  = 0.0 + 0.0*I;
    mgh->fk_tec[index][n]    = mgh->fkm_tec[index][n]  = 0.0 + 0.0*I;
    mgh->fk_lec[index][n]    = mgh->fkm_lec[index][n]  = 0.0 + 0.0*I;
    mgh->fk_cdens[index][n]  = mgh->fkm_cdens[index][n]= 0.0 + 0.0*I;

    for ( m=0; m<sys.molptr->num_mols; m++ ){
        
      kfac   = cexp(I*mgh->k[n]*mols[m].xtrue[1]);
      kfac_m = cexp(-I*mgh->k[n]*mols[m].xtrue[1]);
  
      mgh->fk_scatt[index][n]   +=  kfac;
      mgh->fkm_scatt[index][n]  +=  kfac_m;

      mgh->fk_tv[index][n]   +=  mols[m].v[0]*kfac;
      mgh->fkm_tv[index][n]  +=  mols[m].v[0]*kfac_m;

      mgh->fk_lv[index][n]   +=  mols[m].v[1]*kfac;
      mgh->fkm_lv[index][n]  +=  mols[m].v[1]*kfac_m;

      mgh->fk_tav[index][n]   +=  mols[m].w[2]*kfac;
      mgh->fkm_tav[index][n]  +=  mols[m].w[2]*kfac_m;
     
      mgh->fk_lav[index][n]   +=  mols[m].w[1]*kfac;
      mgh->fkm_lav[index][n]  +=  mols[m].w[1]*kfac_m;

      mgh->fk_tec[index][n]  +=  molcharge[m]*mols[m].v[0]*kfac;
      mgh->fkm_tec[index][n] +=  molcharge[m]*mols[m].v[0]*kfac_m;

      mgh->fk_lec[index][n]  +=  molcharge[m]*mols[m].v[1]*kfac;
      mgh->fkm_lec[index][n] +=  molcharge[m]*mols[m].v[1]*kfac_m;

      mgh->fk_cdens[index][n] += molcharge[m]*kfac;
      mgh->fkm_cdens[index][n] += molcharge[m]*kfac_m;

      mgh->fk_vav[index][n]   +=  mols[m].w[2]*kfac;
      mgh->fkm_vav[index][n]  +=  mols[m].v[0]*kfac_m;

    }
    
    mgh->fk_tv[index][n]  *=  mass/mdens;
    mgh->fkm_tv[index][n] *=  mass/mdens;
    
    mgh->fk_lv[index][n]  *=  mass/mdens;
    mgh->fkm_lv[index][n] *=  mass/mdens;

    mgh->fk_tav[index][n]  *=  3.0*mass/(2.0*mdens);
    mgh->fkm_tav[index][n] *=  3.0*mass/(2.0*mdens);
    
    mgh->fk_lav[index][n]  *=  3.0*mass/(2.0*mdens);
    mgh->fkm_lav[index][n] *=  3.0*mass/(2.0*mdens);

    // Only valid for single-component systems
    mgh->fk_vav[index][n]  *=  3.0*mass/(2.0*mdens);
    mgh->fkm_vav[index][n] *=  3.0*mass/(2.0*mdens);                  
  }

  (mgh->i)++;

  if ( mgh->i == mgh->lvec ){
    for ( unsigned k=0; k<mgh->nwave; k++ ){
      for ( unsigned n=0; n<mgh->lvec; n++ ){
	for ( unsigned nn=0; nn<mgh->lvec-n; nn++ ){
	  mgh->c_tv[n][k]    += mgh->fkm_tv[nn][k]*mgh->fk_tv[n+nn][k];
	  mgh->c_lv[n][k]    += mgh->fkm_lv[nn][k]*mgh->fk_lv[n+nn][k];
	  mgh->c_scatt[n][k] += mgh->fk_scatt[nn][k]*mgh->fkm_scatt[n+nn][k];
	  mgh->c_tav[n][k]   += mgh->fkm_tav[nn][k]*mgh->fk_tav[n+nn][k];
	  mgh->c_lav[n][k]   += mgh->fkm_lav[nn][k]*mgh->fk_lav[n+nn][k];
  	  mgh->c_vav[n][k]   += mgh->fkm_vav[nn][k]*mgh->fk_vav[n+nn][k];

  	  mgh->c_tec[n][k]   += mgh->fkm_tec[nn][k]*mgh->fk_tec[n+nn][k];
  	  mgh->c_lec[n][k]   += mgh->fkm_lec[nn][k]*mgh->fk_lec[n+nn][k];
  	  mgh->c_cdens[n][k] += mgh->fkm_cdens[nn][k]*mgh->fk_cdens[n+nn][k];

	  mgh->c_tcav[n][k]   += mgh->fkm_tec[nn][k]*mgh->fk_tav[n+nn][k];
	  mgh->c_lcav[n][k]   += mgh->fkm_lec[nn][k]*mgh->fk_lav[n+nn][k];

	  mgh->c_cdensav[n][k]   += mgh->fkm_cdens[nn][k]*mgh->fk_tav[n+nn][k];

	}
      }
    }

    (mgh->nsample)++;


    FILE *fout_tv = fopen("mtvacf.dat", "w");
    FILE *fout_lv = fopen("mlvacf.dat", "w");
    FILE *fout_scatt = fopen("mscatt.dat", "w");
    FILE *fout_tav = fopen("mtavacf.dat", "w");
    FILE *fout_lav = fopen("mlavacf.dat", "w");
    FILE *fout_vav = fopen("mvavccf.dat", "w");

    FILE *fout_tec = fopen("mtcacf.dat", "w");
    FILE *fout_lec = fopen("mlcacf.dat", "w");
    FILE *fout_cdens = fopen("mcdensacf.dat", "w");
    FILE *fout_tcav = fopen("mtcavccf.dat", "w");
    FILE *fout_lcav = fopen("mlcavccf.dat", "w");
    
    FILE *fout_cdensav = fopen("mcdensavccf.dat", "w");

    for ( unsigned n=0; n<mgh->lvec; n++ ){
      double t   = (n+0.5)*mgh->dtsample;
      
      fprintf(fout_tv, "%f ", t);
      fprintf(fout_lv, "%f ", t);
      fprintf(fout_scatt, "%f ", t);
      fprintf(fout_tav, "%f ", t);
      fprintf(fout_lav, "%f ", t);
      fprintf(fout_vav, "%f ", t);

      fprintf(fout_tec, "%f ", t);
      fprintf(fout_lec, "%f ", t);
      fprintf(fout_cdens, "%f ", t);
      fprintf(fout_tcav, "%f ", t);
      fprintf(fout_lcav, "%f ", t);

      fprintf(fout_cdensav, "%f ", t);

      double fac = 1.0/(mgh->nsample*sys.volume*(mgh->lvec-n));
      
      for ( unsigned k=0; k<mgh->nwave; k++ ){
	fprintf(fout_tv, "%f %f ", creal(mgh->c_tv[n][k]*fac), 
		cimag(mgh->c_tv[n][k]*fac));
	fprintf(fout_lv, "%f %f ", creal(mgh->c_lv[n][k]*fac), 
		cimag(mgh->c_lv[n][k]*fac));
	fprintf(fout_scatt, "%f %f ", creal(mgh->c_scatt[n][k]*fac), 
		cimag(mgh->c_scatt[n][k]*fac));
	fprintf(fout_tav, "%f %f ", creal(mgh->c_tav[n][k]*fac), 
		cimag(mgh->c_tav[n][k]*fac));
	fprintf(fout_lav, "%f %f ", creal(mgh->c_lav[n][k]*fac), 
		cimag(mgh->c_lav[n][k]*fac));
	fprintf(fout_vav, "%f %f ", 
		creal(mgh->c_vav[n][k]*fac), cimag(mgh->c_vav[n][k]*fac));

	fprintf(fout_tec, "%f %f ", 
		creal(mgh->c_tec[n][k]*fac), cimag(mgh->c_tec[n][k]*fac));
	fprintf(fout_lec, "%f %f ", 
		creal(mgh->c_lec[n][k]*fac), cimag(mgh->c_lec[n][k]*fac));
	fprintf(fout_cdens, "%f %f ", 
		creal(mgh->c_cdens[n][k]*fac), cimag(mgh->c_cdens[n][k]*fac));
	fprintf(fout_tcav, "%f %f ", 
		creal(mgh->c_tcav[n][k]*fac), cimag(mgh->c_tcav[n][k]*fac));
	fprintf(fout_lcav, "%f %f ", 
		creal(mgh->c_lcav[n][k]*fac), cimag(mgh->c_lcav[n][k]*fac));

	fprintf(fout_cdensav, "%f %f ", 
		creal(mgh->c_cdensav[n][k]*fac), cimag(mgh->c_cdensav[n][k]*fac));

      }
      fprintf(fout_tv, "\n"); 
      fprintf(fout_lv, "\n"); 
      fprintf(fout_scatt, "\n");
      fprintf(fout_tav, "\n"); 
      fprintf(fout_lav, "\n"); 
      fprintf(fout_vav, "\n"); 

      fprintf(fout_tec, "\n"); 
      fprintf(fout_lec, "\n"); 
      fprintf(fout_cdens, "\n"); 
      fprintf(fout_tcav, "\n"); 
      fprintf(fout_lcav, "\n"); 

      fprintf(fout_cdensav, "\n"); 
    }
	
    fclose(fout_tv);
    fclose(fout_lv);
    fclose(fout_scatt);
    fclose(fout_tav);
    fclose(fout_lav);
    fclose(fout_vav);

    fclose(fout_tec);
    fclose(fout_lec);    
    fclose(fout_cdens);
    fclose(fout_tcav);
    fclose(fout_lcav);

    fclose(fout_cdensav);
    mgh->i = 0;
  }

  free(molcharge);
}



// Hydro profiler (atomic) 
sepprofs *sep_profs_init(char type, int lvec, int isample){
  
  sepprofs *ptr = malloc(sizeof(sepprofs));
  if ( ptr==NULL )
    sep_error("%s at line %d: Couldn't allocate memory", __func__, __LINE__);
  
  ptr->lvec = lvec;
  ptr->nsample = 0;
  ptr->isample = isample;
  ptr->type = type;

  // default - can be changed 
  ptr->dir = 2;
  ptr->dirvel = 0;

  ptr->momc  = sep_vector(lvec);
  ptr->svel  = sep_vector(lvec);
  ptr->dens  = sep_vector(lvec);
  ptr->temp  = sep_vector(lvec);

  return ptr;
}

void sep_profs_close(sepprofs *ptr){
 

  free(ptr->momc);  free(ptr->svel);
  free(ptr->dens);  free(ptr->temp);

}



void sep_profs_sampler(seppart *pptr, sepprofs *ptr, sepsys sys){
  double dl, *j, *rho, *sumv2, dV;
  int n, i, k=0, *numb;
  FILE *fout;

  const int dir = ptr->dir;
  const int dirvel = ptr->dirvel;

  dl = sys.length[dir]/ptr->lvec;

  if ( dir == 0 )
    dV = dl*sys.length[1]*sys.length[2];
  else if ( dir == 1 )
    dV = sys.length[0]*dl*sys.length[2];
  else 
    dV = sys.length[0]*sys.length[1]*dl;
  	     		
    
  j = sep_vector(ptr->lvec);
  rho = sep_vector(ptr->lvec);
  sumv2 = sep_vector(ptr->lvec);
  numb = sep_vector_int(ptr->lvec);

  
  for ( n=0; n<sys.npart; n++ ){
    if ( pptr[n].type == ptr->type ) {
      
      i = pptr[n].x[dir]/dl;            
      j[i] += pptr[n].m*pptr[n].v[dirvel];
      rho[i] += pptr[n].m;

      for ( k=0; k<3; k++ ){			
	if ( k != dirvel ) 
	  sumv2[i] += pptr[n].m*sep_Sq(pptr[n].v[k]);			
      }

      numb[i]++;
    }
  }
  
  (ptr->nsample)++;
  double idV = 1.0/dV;

  for ( unsigned n=0; n<ptr->lvec; n++ ){
    ptr->momc[n] += j[n]*idV;
    ptr->dens[n] += rho[n]*idV;
    if ( numb[n] > 0 )
      ptr->temp[n] += sumv2[n]; 
  }
  
  free(j); free(rho); free(sumv2); free(numb);
  
  
  if ( ptr->nsample%100 == 0 ){

    fout = fopen("profs.dat", "w");
    if ( fout == NULL )
      sep_error("%s at %d: Couldn't open file", __func__, __LINE__);
    
    for ( unsigned n=0; n<ptr->lvec; n++ ){

      double insample = 1.0/ptr->nsample;
      double temp_fac = 0;
      if ( ptr->dens[n] > 0.0 ){
	ptr->svel[n] = ptr->momc[n]/ptr->dens[n];	
	temp_fac = 1.0/(2.0*dV*ptr->dens[n]);
      }     
      
      double z = (n + 0.5)*dl;
      fprintf(fout, "%f %f %f %f %f\n",  z, ptr->momc[n]*insample, 
	      ptr->dens[n]*insample, ptr->temp[n]*temp_fac, ptr->svel[n]); 
	
    }
    fclose(fout);
  }  

}				


// Hydro profiler (molecular) 
sepmprofs *sep_mprofs_init(char type, int lvec, int isample, int dir, int dirvel, 
			   int diramom){
  
  sepmprofs *ptr = malloc(sizeof(sepprofs));
  if ( ptr==NULL )
    sep_error("%s at line %d: Couldn't allocate memory", __func__, __LINE__);

  ptr->lvec    = lvec;
  ptr->nsample = 0;
  ptr->isample = isample;
  ptr->type    = type;
  ptr->dir     = dir;
  ptr->dirvel  = dirvel;
  ptr->diramom = diramom;

  ptr->momc  = sep_vector(lvec);
  ptr->dens  = sep_vector(lvec);
  ptr->amom  = sep_matrix(lvec, 3);
  ptr->temp  = sep_vector(lvec);
  ptr->numb  = sep_vector(lvec);

  ptr->inertia = sep_tensor(lvec,3,3);

  return ptr;
}

void sep_mprofs_close(sepmprofs *ptr){
 
  
  free(ptr->momc);  
  sep_free_matrix(ptr->amom, ptr->lvec);
  free(ptr->dens);  
  free(ptr->temp);
  free(ptr->numb);

  sep_free_tensor(ptr->inertia, ptr->lvec, 3);
  
}



void sep_mprofs_sampler(seppart *pptr, sepmol *mols, 
			sepmprofs *ptr, sepsys sys){
  double dl, dV, w[3]={0.0};
  unsigned int n, i, k=0;
  FILE *fout;

  sep_mol_spin(pptr, mols, sys, 0); //cm is also evaluated here
  sep_mol_velcm(pptr, mols, sys);

  const int dir     = ptr->dir;
  const int dirvel  = ptr->dirvel;
  const int diramom = ptr->diramom;

  dl = sys.length[dir]/ptr->lvec;

  if ( dir == 0 )
    dV = dl*sys.length[1]*sys.length[2];
  else if ( dir == 1 )
    dV = sys.length[0]*dl*sys.length[2];
  else 
    dV = sys.length[0]*sys.length[1]*dl;
  
  for ( n=0; n<sys.molptr->num_mols; n++ ){
    if ( mols[n].type == ptr->type ) {
      
      i = mols[n].x[dir]/dl;            
      
      if ( i >= ptr->lvec ) 
	sep_error("%s at %d: Array index larger than vector length", 
		  __func__, __LINE__);

      ptr->momc[i] += mols[n].m*mols[n].v[dirvel];
      ptr->dens[i] += mols[n].m;
      ptr->numb[i] ++;
      for ( k=0; k<3; k++ ) ptr->amom[i][k] += mols[n].s[k];
      
      for ( k=0; k<3; k++ ){			
	if ( k != (unsigned)dirvel ) 
	  ptr->temp[i] += mols[n].m*sep_Sq(mols[n].v[k]); 		
	for ( int kk=0; kk<3; kk++ )
	  ptr->inertia[i][k][kk] += mols[n].inertia[k][kk];
      }
      
    }
  }
  
  (ptr->nsample)++;

  if ( ptr->nsample%100 == 0 ){
    
    fout = fopen("mprofs.dat", "w");
    if ( fout == NULL )
      sep_error("%s at %d: Couldn't open file", __func__, __LINE__);

    double fac = 1/(ptr->nsample*dV);    
    for ( unsigned n=0; n<ptr->lvec; n++ ){

      double **inertia = sep_matrix(3,3);

      double x = (n + 0.5)*dl;

      if ( ptr->numb[n] > 0 ){
	
	double a = ptr->dens[n]*fac;          // Mass density 
	double b = ptr->momc[n]*fac;          // Momentum current density   
	
	double av_numb = ptr->numb[n]/ptr->nsample; 
	double c = ptr->temp[n]/(2*ptr->nsample*av_numb); //Temperature

	double spin[3];
	for (  int k=0; k<3; k++ ){
	  spin[k] = ptr->amom[n][k]*fac;
	  for ( int kk=0; kk<3; kk++ )
	    inertia[k][kk] = ptr->inertia[n][k][kk]*fac; //(a*ptr->nsample);
	}

	sep_solvelineq1(w, spin, inertia, 3);
	fprintf(fout, "%e %e %e %e %e ",  x, a, b/a, c, w[diramom]);
	fprintf(fout, "%e %e %e %e %e %e %e %e %e %e %e %e\n", 
		spin[0], spin[1], spin[2], 
		inertia[0][0], inertia[0][1], inertia[0][2], 
		inertia[1][0], inertia[1][1], inertia[1][2],
		inertia[2][0], inertia[2][1], inertia[2][2]);
	
      }
      else {
	fprintf(fout, "%e %e %e %e %e %e ",  x, 0.0, 0.0, 0.0, 0.0, 0.0);
	fprintf(fout, "%e %e % e%e %e %e %e %e %e %e %e %e\n", 
		0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0);
      }
      sep_free_matrix(inertia, 3);

      
    }
    fclose(fout);
  }  
  
}				


// Molecular couple tensor ACFs
sepmcacf *sep_mcacf_init(int lvec, double tsample, double dt){
  
  sepmcacf *ptr = malloc(sizeof(sepmcacf));
  if ( ptr==NULL )
    sep_error("%s at line %d: Couldn't allocate memory", __func__, __LINE__);
  
  ptr->dt   = dt;
  ptr->lvec = lvec;
  ptr->i    = 0; 
  
  ptr->nsample = 0;

  ptr->dtsample = tsample/lvec;
  ptr->isample = (int)(ptr->dtsample/dt);
  if ( ptr->isample < 1 ) 
    sep_error("%s at line %d: isample is too small - CHECK lvec argument",
	      __func__, __LINE__);


  ptr->cacf = sep_matrix(lvec, 2);
  ptr->scouple = sep_matrix(lvec, 3);
  ptr->acouple = sep_matrix(lvec, 3);

  return ptr;
}


void sep_mcacf_sample(sepmcacf *sptr, sepatom *atoms, sepmol *mols, 
		      sepret *ret, sepsys sys){

  sep_eval_couple_tensor(atoms, mols, ret, sys);

  int index = sptr->i;

  sptr->scouple[index][0] = 0.5*(ret->T_mol[0][1]+ret->T_mol[1][0]);
  sptr->scouple[index][1] = 0.5*(ret->T_mol[0][2]+ret->T_mol[2][0]);
  sptr->scouple[index][2] = 0.5*(ret->T_mol[1][2]+ret->T_mol[2][1]);

  sptr->acouple[index][0] = 0.5*(ret->T_mol[0][1]-ret->T_mol[1][0]);
  sptr->acouple[index][1] = 0.5*(ret->T_mol[0][2]-ret->T_mol[2][0]);
  sptr->acouple[index][2] = 0.5*(ret->T_mol[1][2]-ret->T_mol[2][1]);

  (sptr->i)++;

  if ( sptr->i == sptr->lvec ){
     
    for ( int k=0; k<3; k++ ){
      for ( unsigned n=0; n<sptr->lvec; n++ ){
        for ( unsigned nn=0; nn<sptr->lvec-n; nn++ ){
          sptr->cacf[n][0] += sptr->scouple[nn][k]*sptr->scouple[n+nn][k];
	  sptr->cacf[n][1] += sptr->acouple[nn][k]*sptr->acouple[n+nn][k];
        }
      }
    }

    (sptr->nsample)++;

    FILE *fout1 = fopen("mcacf.dat", "w");
   
    for ( unsigned n=0; n<sptr->lvec; n++ ){
      double t   = (n+0.5)*sptr->dtsample;
      double fac = sys.volume/(3*(sptr->lvec-n)*sptr->nsample);
      fprintf(fout1, "%f %f %f\n", t, sptr->cacf[n][0]*fac, sptr->cacf[n][1]*fac);
    }

    fclose(fout1); 
    
    sptr->i = 0;
    
  }

}

void sep_mcacf_close(sepmcacf *ptr){
  
 
  sep_free_matrix(ptr->cacf, ptr->lvec);
  sep_free_matrix(ptr->scouple, ptr->lvec);
  sep_free_matrix(ptr->acouple, ptr->lvec);
  
  free(ptr);

}

// Molecular velocity ACF
sepmvacf *sep_mvacf_init(int lvec, double tsample, sepsys sys){
  
  sepmvacf *ptr = malloc(sizeof(sepmvacf));
  if ( ptr==NULL )
    sep_error("%s at line %d: Couldn't allocate memory", __func__, __LINE__);
  
  ptr->dt   = sys.dt;
  ptr->lvec = lvec;
  ptr->i    = 0; 
  
  ptr->nsample = 0;
  
  ptr->dtsample = tsample/lvec;
  ptr->isample = (int)(ptr->dtsample/sys.dt);
  if ( ptr->isample < 1 ) 
    sep_error("%s at line %d: isample is too small - CHECK lvec argument",
	      __func__, __LINE__);

  ptr->vacf = sep_vector(lvec);
  ptr->vels = sep_matrix(lvec, sys.molptr->num_mols);
  
  return ptr;
}

 
void sep_mvacf_sample(seppart *ptr, sepmvacf *vptr, sepmol *mols, sepsys sys){
    
  sep_mol_velcm(ptr, mols, sys);

  int index = vptr->i;

  const int num_mols = sys.molptr->num_mols;
  for ( int i=0; i<num_mols; i++ )
    vptr->vels[index][i] = mols[i].v[0];
  
  (vptr->i)++;
  
  if ( vptr->i == vptr->lvec ){
    
    for ( int i=0; i<num_mols; i++ )
      for ( unsigned k=0; k<vptr->lvec; k++ )
        for ( unsigned kk=0; kk<vptr->lvec-k; kk++ )
          vptr->vacf[k] += vptr->vels[kk][i]*vptr->vels[k+kk][i];   
  
    (vptr->nsample)++;

    FILE *fout = fopen("mvacf.dat", "w");
    for ( unsigned k=0; k<vptr->lvec; k++ ){
      double t   = (k+0.5)*vptr->dtsample;
      double fac = 1.0/((vptr->lvec-k)*sys.npart*vptr->nsample);
      fprintf(fout, "%f %f\n", t, vptr->vacf[k]*fac);
    }
    fclose(fout);

    vptr->i=0;
    
  }

}

void sep_mvacf_close(sepmvacf *ptr){

  free(ptr->vacf);
  sep_free_matrix(ptr->vels, ptr->lvec);
 
  free(ptr);

}

// Molecular angular velocity ACF
sepmavacf *sep_mavacf_init(int lvec, double tsample, sepsys sys){
  
  sepmavacf *ptr = malloc(sizeof(sepmavacf));
  if ( ptr==NULL )
    sep_error("%s at line %d: Couldn't allocate memory", __func__, __LINE__);
  
  ptr->dt   = sys.dt;
  ptr->lvec = lvec;
  ptr->i    = 0; 
  
  ptr->nsample = 0;
  
  ptr->dtsample = tsample/lvec;
  ptr->isample = (int)(ptr->dtsample/sys.dt);
  if ( ptr->isample < 1 ) 
    sep_error("%s at line %d: isample is too small - CHECK lvec argument",
	      __func__, __LINE__);

  ptr->avacf = sep_vector(lvec);
  ptr->avels = sep_matrix(lvec, sys.molptr->num_mols);
  
  return ptr;
}

 
void sep_mavacf_sample(seppart *ptr, sepmavacf *vptr, sepmol *mols, sepsys sys){
    
  sep_mol_spin(ptr, mols, sys, 1);


  // To be removed - from here
  sep_calc_mol_ete(ptr, mols, 'A', 0, mols[0].nuau-1, sys);
  double q = fabs(ptr[0].z);
  // to here

  int index = vptr->i;

  // To be removed - from here
  const int num_mols = sys.molptr->num_mols;

  double p[3];
  for ( int i=0; i<num_mols; i++ ){
    double mass = mols[i].m;
    for ( int k=0; k<3; k++ )
      p[k] = mols[i].ete[k]*q*mass;

    //vptr->avels[index][i] = mols[i].w[1]*p[2]-mols[i].w[2]*p[1];
    vptr->avels[index][i] = p[1];
  }
  // To here
	
  //  const int num_mols = sys.molptr->num_mols;
  //  for ( int i=0; i<num_mols; i++ )
    //    vptr->avels[index][i] = mols[i].w[0];
  
  (vptr->i)++;
  
  if ( vptr->i == vptr->lvec ){
    
    for ( int i=0; i<num_mols; i++ )
      for ( unsigned k=0; k<vptr->lvec; k++ )
        for ( unsigned kk=0; kk<vptr->lvec-k; kk++ )
          vptr->avacf[k] += vptr->avels[kk][i]*vptr->avels[k+kk][i];   
  
    (vptr->nsample)++;

    FILE *fout = fopen("mavacf.dat", "w");
    for ( unsigned k=0; k<vptr->lvec; k++ ){
      double t   = (k+0.5)*vptr->dtsample;
      double fac = 1.0/((vptr->lvec-k)*sys.molptr->num_mols*vptr->nsample);
      fprintf(fout, "%f %f\n", t, vptr->avacf[k]*fac);
    }
    fclose(fout);

    vptr->i=0;
    
  }

}

void sep_mavacf_close(sepmavacf *ptr){

  free(ptr->avacf);
  sep_free_matrix(ptr->avels, ptr->lvec);
 
  free(ptr);

}

